import torch
from torch.utils.data import DataLoader
import torch.nn as nn
from datasets.dataset import TicketDataset
from models.MPDnnModel import PricePredictor
from utils.preprocess import preprocess
from utils.concat_csv_files import concat_csv_files
from tqdm import tqdm
import json
import matplotlib.pyplot as plt
from sklearn.metrics import r2_score, mean_absolute_percentage_error
import numpy as np

def evaluate(model_path, test_df, batch_size=32):
    print(f"\n🔎 正在评估模型：{model_path}")

    # 加载模型结构参数
    with open("checkpoints/model_config.json", "r") as f:
        config = json.load(f)

    model = PricePredictor(
        num_numeric_features=config["num_numeric_features"],
        num_film_ids=config["num_film_ids"],
        film_emb_dim=config["film_emb_dim"],
        hidden_dims=config["hidden_dims"],
        dropout=config["dropout"]
    )

    model.load_state_dict(torch.load(model_path))
    model.eval()

    test_dataset = TicketDataset(test_df)
    test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)

    criterion = nn.MSELoss()
    mae_loss = nn.L1Loss()

    total_mse = 0.0
    total_mae = 0.0
    preds = []
    trues = []

    with torch.no_grad():
        for x_numeric, x_film_id, y in tqdm(test_loader, desc="🧪 Testing"):
            output = model(x_numeric, x_film_id)
            loss_mse = criterion(output, y)
            loss_mae = mae_loss(output, y)

            total_mse += loss_mse.item() * x_numeric.size(0)
            total_mae += loss_mae.item() * x_numeric.size(0)

            preds.extend(output.squeeze().tolist())
            trues.extend(y.squeeze().tolist())

    preds = np.array(preds)
    trues = np.array(trues)

    mse = total_mse / len(test_loader.dataset)
    mae = total_mae / len(test_loader.dataset)
    r2 = r2_score(trues, preds)
    mape = mean_absolute_percentage_error(trues, preds)
    acc = np.mean(np.abs(preds - trues) / (trues + 1e-8) < 0.01)

    print(f"📈 测试集 MSE: {mse:.4f} | MAE: {mae:.4f} | R^2: {r2:.4f} | MAPE: {mape:.4%} | ACC(<1%): {acc:.4%}")

    # 中文图像可视化
    plt.figure(figsize=(6, 5))
    plt.scatter(trues, preds, alpha=0.5)
    plt.xlabel("真实价格")
    plt.ylabel("预测价格")
    plt.title("真实值 vs 预测值")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("outputs/真实值_vs_预测值.png")
    plt.close()

    plt.figure(figsize=(6, 5))
    residuals = preds - trues
    plt.scatter(trues, residuals, alpha=0.5)
    plt.axhline(0, color='red', linestyle='--')
    plt.xlabel("真实价格")
    plt.ylabel("残差")
    plt.title("残差图")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("outputs/残差图.png")
    plt.close()

    plt.figure(figsize=(6, 5))
    plt.hist(residuals, bins=50, alpha=0.7)
    plt.xlabel("预测误差")
    plt.ylabel("频率")
    plt.title("误差分布图")
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("outputs/误差分布图.png")
    plt.close()

    # 新增对比图：预测值与真实值误差折线
    plt.figure(figsize=(10, 4))
    plt.plot(preds[:2000], label='预测值')
    plt.plot(trues[:2000], label='真实值')
    plt.title("预测值与真实值前100个样本对比")
    plt.xlabel("样本编号")
    plt.ylabel("价格")
    plt.legend()
    plt.grid(True)
    plt.tight_layout()
    plt.savefig("outputs/预测值_真实值_对比前100.png")
    plt.close()

    return preds.tolist(), trues.tolist()


if __name__ == "__main__":
    TRAIN_DIR = "data/train"
    TEST_DIR = "data/test"

    # 重新读取并预处理原始测试集
    _, test_raw = concat_csv_files(TRAIN_DIR), concat_csv_files(TEST_DIR)
    _, test_df = preprocess(_, test_raw, save_dir="data")

    # 调用测试函数
    evaluate("checkpoints/best_model.pt", test_df)
